import time
import secrets
import hashlib

# 导入你的加密算法实现
from rsa_core import RSA
from elgamal_core import ElGamal
# 修改: 导入 get_curve_by_name 函数，不再导入 secp256k1_curve
from ecc_core import (
    get_curve_by_name,
    generate_keypair as ecc_gen_keys,
    ecc_encrypt_simplified,
    ecc_decrypt_simplified,
)
from math_utils import power, extended_gcd, mod_inverse

# ==============================================================================
# Part 1: 选择密文攻击 (Chosen-Ciphertext Attack) 演示
# ==============================================================================

def test_rsa_cca_resistance():
    """测试: RSA with OAEP 对选择密文攻击的抵抗性。"""
    print("\n" + "="*20, "测试: RSA-OAEP CCA 抵抗性", "="*20)
    rsa_instance = RSA(bits=1024)
    message = b"This is the original secret message for RSA."
    n, e = rsa_instance.public_key
    ciphertext_int = rsa_instance.encrypt(message)
    print("生成了一个合法的RSA密文。")
    tampered_ciphertext_int = (ciphertext_int * power(2, e, n)) % n
    print("攻击者将密文在数学上进行了篡改。")
    try:
        rsa_instance.decrypt(tampered_ciphertext_int)
        print(f"!! 失败：被篡改的密文竟然成功解密了。")
    except ValueError as err:
        print(f"++ 成功：解密被篡改的密文时，OAEP填充校验失败，抛出错误: {err}")
        print("结论: RSA-OAEP 成功抵抗了这种形式的CCA。")

def test_elgamal_malleability_cca():
    """演示: "教科书式"ElGamal 的延展性 (Malleability) 漏洞。"""
    print("\n" + "="*20, "演示: ElGamal 延展性 (CCA)", "="*20)
    elgamal_instance = ElGamal(key_size_bits=1024)
    p, _, g, _ = elgamal_instance.public_key
    original_amount = 100
    print(f"原始消息 (交易金额): M = {original_amount}")
    c1, c2 = ElGamal.encrypt(elgamal_instance.public_key, original_amount)
    print(f"为 M 生成的合法密文: (c1, c2)")
    multiplier = 10
    tampered_c2 = (c2 * multiplier) % p
    tampered_ciphertext = (c1, tampered_c2)
    print(f"攻击者将 c2 乘以 {multiplier}，构造了新的密文。")
    public_params = (p, None, g)
    decrypted_tampered_amount = ElGamal.decrypt(elgamal_instance.private_key, public_params, tampered_ciphertext)
    print(f"被篡改的密文解密后的结果: M' = {decrypted_tampered_amount}")
    if decrypted_tampered_amount == original_amount * multiplier:
        print(f"++ 成功演示漏洞: 解密结果 ({decrypted_tampered_amount}) 正是 原始金额({original_amount}) * 乘数({multiplier})。")
        print("结论: 教科书式ElGamal是可延展的，不具备CCA安全性。")
    else:
        print("!! 演示失败，出现未知错误。")

def test_ecc_cca_resistance():
    """测试: ECC 混合加密方案对选择密文攻击的抵抗性。"""
    print("\n" + "="*20, "测试: ECC CCA 抵抗性", "="*20)
    # 修改: 使用 get_curve_by_name 获取曲线实例
    curve = get_curve_by_name("secp256k1")
    priv_key, pub_key = ecc_gen_keys(curve)
    message = b"This is the original secret message for ECC."
    R, C = ecc_encrypt_simplified(pub_key, curve, message)
    print("生成了一个合法的ECC密文 (R, C)。")
    tampered_R = curve.add_points(R, curve.G)
    print("攻击者篡改了密文的 R 点。")
    decrypted_tampered = ecc_decrypt_simplified(priv_key, curve, tampered_R, C)
    print(f"原始消息: {message}")
    print(f"解密篡改后的密文得到: {decrypted_tampered}")
    if message != decrypted_tampered:
        print("++ 成功: 解密被篡改的密文得到了无意义的乱码。")
        print("结论: 这种混合加密方案能抵抗简单的延展性攻击。")
    else:
        print("!! 失败: 被篡改的密文竟然解密回了原始消息。")

# ... (Part 2 and other functions from previous response) ...
def _get_key_with_specific_weight(length_bits, target_weight):
    if target_weight > length_bits: raise ValueError("权重不能大于位长")
    bits = [1] * target_weight + [0] * (length_bits - target_weight)
    secrets.SystemRandom().shuffle(bits)
    if bits[0] == 0:
        for i in range(1, length_bits):
            if bits[i] == 1:
                bits[0], bits[i] = bits[i], bits[0]
                break
    return int("".join(map(str, bits)), 2)

def simulate_timing_attack_vulnerability():
    print("\n" + "="*20, "模拟: 时序攻击 (Timing Attack) 漏洞", "="*20)
    print("本实验并非真实攻击，仅为证明漏洞存在。")
    key_size, num_iterations = 512, 200
    key_low = _get_key_with_specific_weight(key_size, key_size // 4)
    key_high = _get_key_with_specific_weight(key_size, key_size * 3 // 4)
    print(f"低/高权重密钥已生成。")
    modulus = secrets.randbits(key_size) | (1 << (key_size - 1))
    messages = [secrets.randbits(key_size) for _ in range(num_iterations)]
    
    start_time = time.perf_counter()
    for msg in messages: power(msg, key_low, modulus)
    time_low = time.perf_counter() - start_time
    print(f"使用低权重密钥进行 {num_iterations} 次操作的总时间: {time_low:.6f} 秒")
    
    start_time = time.perf_counter()
    for msg in messages: power(msg, key_high, modulus)
    time_high = time.perf_counter() - start_time
    print(f"使用高权重密钥进行 {num_iterations} 次操作的总时间: {time_high:.6f} 秒")
    
    if time_high > time_low:
        diff = (time_high - time_low) / time_low * 100
        print(f"++ 成功观察到时序差异: 高权重密钥的操作时间比低权重长 {diff:.2f}%。")
        print("结论: power() 函数的性能依赖于密钥，存在时序攻击漏洞。")
    else:
        print("!! 未观察到预期的时序差异（可能由系统噪音导致），但算法层面的漏洞依然存在。")

def test_ecc_nonce_reuse_attack():
    print("\n" + "="*20, "演示: ECC K值重用攻击", "="*20)
    # 修改: 使用 get_curve_by_name 获取曲线实例
    curve = get_curve_by_name("secp256k1")
    priv_key, pub_key = ecc_gen_keys(curve)
    reused_k_e = secrets.randbelow(curve.n - 1) + 1
    print(f"Alice不幸地选择了一个k_e并打算重用它。")

    def insecure_ecc_encrypt(k, public_key, curve, message):
        R = curve.scalar_multiply(k, curve.G)
        S_point = curve.scalar_multiply(k, public_key)
        shared_secret_material = S_point[0].to_bytes((S_point[0].bit_length() + 7) // 8, 'big')
        symmetric_key = hashlib.sha256(shared_secret_material).digest()
        return R, bytes(c ^ symmetric_key[i % len(symmetric_key)] for i, c in enumerate(message))

    msg1 = b"Meeting at 10 AM at the docks."
    R1, C1 = insecure_ecc_encrypt(reused_k_e, pub_key, curve, msg1)
    msg2 = b"The password for the server is: StrongP@ssw0rd"
    R2, C2 = insecure_ecc_encrypt(reused_k_e, pub_key, curve, msg2)
    print(f"消息1: {msg1.decode()}")
    print(f"消息2: {msg2.decode()[:20]}...")
    print(f"因为k_e重用, R1 和 R2 必然相同: {R1 == R2}")
    
    m1_xor_m2 = bytes(b1 ^ b2 for b1, b2 in zip(C1, C2))
    guessed_msg1 = b"Meeting at 10 AM at the docks."
    recovered_msg2 = bytes(b1 ^ b2 for b1, b2 in zip(m1_xor_m2, guessed_msg1))
    print(f"Eve基于对消息1的猜测，恢复出消息2: {recovered_msg2.decode()}")
    if recovered_msg2 == msg2:
        print("++ 成功演示攻击: Eve成功恢复了第二条消息的完整内容！")
        print("结论: K值(Nonce)重用是致命的。")

def test_rsa_common_modulus_attack():
    print("\n" + "="*20, "演示: RSA 共模攻击", "="*20)
    rsa_user1 = RSA(bits=1024)
    n = rsa_user1.public_key[0]
    e1 = rsa_user1.public_key[1]
    e2 = 65539
    print(f"用户1的公钥: (n, e1={e1})")
    print(f"用户2的公钥: (n, e2={e2})")
    
    message = b"This is a top secret message for common modulus attack."
    m_int = int.from_bytes(message, 'big')
    c1 = power(m_int, e1, n)
    c2 = power(m_int, e2, n)
    print("同一个明文M，被分别用e1和e2加密，得到C1和C2。")
    
    gcd, a, b = extended_gcd(e1, e2)
    if gcd != 1: print("!! 攻击前提不成立: e1和e2不互素。"); return

    if a < 0: term1 = power(mod_inverse(c1, n), -a, n)
    else: term1 = power(c1, a, n)
    if b < 0: term2 = power(mod_inverse(c2, n), -b, n)
    else: term2 = power(c2, b, n)
        
    recovered_m_int = (term1 * term2) % n
    recovered_message = recovered_m_int.to_bytes((recovered_m_int.bit_length() + 7) // 8, 'big')
    print(f"恢复出的明文: {recovered_message.decode()}")
    if recovered_message == message:
        print("++ 成功演示攻击: 攻击者在不知道任何私钥的情况下恢复了明文！")
        print("结论: 绝不能在不同用户间共享RSA模数n。")

def power_with_fault_injection(base, exp, mod, fault_chance=0.0):
    res = 1; base %= mod
    fault_injected = False
    if secrets.SystemRandom().random() < fault_chance:
        fault_injected = True
        bit_to_flip = secrets.randbelow(base.bit_length())
        base ^= (1 << bit_to_flip)
    while exp > 0:
        if exp % 2 == 1: res = (res * base) % mod
        base = (base * base) % mod
        exp //= 2
    return res, fault_injected

def test_fault_injection_simulation():
    print("\n" + "="*20, "演示: 故障注入攻击模拟", "="*20)
    rsa_instance = RSA(bits=1024)
    message = b"A message to test fault attacks."
    ciphertext = rsa_instance.encrypt(message)
    normal_decrypted = rsa_instance.decrypt(ciphertext)
    print(f"正常解密结果: {normal_decrypted}")
    print("现在模拟在解密计算中注入一个瞬时故障...")
    n, d = rsa_instance.private_key
    m_int_faulty, fault_was_injected = power_with_fault_injection(ciphertext, d, n, fault_chance=1.0)
    if fault_was_injected:
        print("故障已成功注入到解密计算中。")
        try:
            k = (n.bit_length() + 7) // 8
            em_faulty = m_int_faulty.to_bytes(k, 'big')
            decrypted_faulty = rsa_instance._oaep_decode(em_faulty, k, b"")
            print(f"故障解密后的（可能是乱码的）结果: {decrypted_faulty}")
        except ValueError as e:
            print(f"故障解密导致OAEP解码失败: {e}")
        print("++ 成功演示: 一个计算上的小错误导致了解密结果的巨大差异。")
        print("结论: 物理层面的故障注入是严重威胁，需要硬件级防护。")

# textbook_rsa_cca_attack.py

import secrets
from publickey.math_utils import generate_prime_number, extended_gcd, mod_inverse, power

class TextbookRSA:
    """
    一个不使用任何填充的“教科书式”RSA实现，用于演示漏洞。
    警告：绝对不能在实际应用中使用此类实现！
    """
    def __init__(self, bits=1024):
        p = generate_prime_number(bits // 2)
        q = generate_prime_number(bits // 2)
        while p == q:
            q = generate_prime_number(bits // 2)
        
        self.n = p * q
        phi_n = (p - 1) * (q - 1)
        self.e = 65537
        while extended_gcd(self.e, phi_n)[0] != 1:
            self.e = 2 * secrets.randbelow((phi_n - 3) // 2) + 3 # 备用e
        
        self.d = mod_inverse(self.e, phi_n)

    def encrypt(self, message_int: int) -> int:
        """加密一个整数"""
        if not (0 <= message_int < self.n):
            raise ValueError("消息必须在 [0, n-1) 范围内")
        return power(message_int, self.e, self.n)

    def decrypt(self, ciphertext_int: int) -> int:
        """解密一个整数"""
        return power(ciphertext_int, self.d, self.n)

def demonstrate_textbook_rsa_cca():
    """
    主函数：完整演示对教科书式RSA的选择密文攻击。
    """
    print("="*20, "演示对教科书式RSA的CCA攻击", "="*20)

    # --- 场景设置 ---
    # 1. 受害者(Victim)创建一个教科书式RSA密钥对
    victim_rsa = TextbookRSA(bits=1024)
    print("受害者创建了一个不安全的、无填充的RSA密钥对。")
    
    # 2. 原始的秘密消息
    original_message = b"Your final exam grade is: A+"
    original_message_int = int.from_bytes(original_message, 'big')
    print(f"原始的秘密消息: '{original_message.decode()}' (整数形式: {original_message_int})")

    # 3. 受害者加密消息，准备发送
    # (在真实场景中，这一步由发送方完成，密文被攻击者截获)
    ciphertext_c = victim_rsa.encrypt(original_message_int)
    print(f"原始密文 C 被攻击者截获。")

    print("\n" + "-"*15, "攻击开始", "-"*15)

    # --- 攻击者 (Attacker) 的操作 ---
    # 4. 攻击者选择一个随机数 r
    r = secrets.randbelow(victim_rsa.n - 2) + 2 # 选择一个 [2, n-1] 范围的 r
    print(f"1. 攻击者选择一个随机数 r = {r}")

    # 5. 攻击者计算 r^e mod n
    r_pow_e = power(r, victim_rsa.e, victim_rsa.n)

    # 6. 攻击者构造新的、篡改过的密文 C'
    tampered_ciphertext_c_prime = (ciphertext_c * r_pow_e) % victim_rsa.n
    print(f"2. 攻击者构造了新的密文 C' = C * (r^e) mod n")

    # --- “解密预言机” (Decryption Oracle) 交互 ---
    # 7. 攻击者将 C' 发送给受害者请求解密
    print(f"3. 攻击者将 C' 发送给受害者的解密服务（预言机）。")
    decrypted_m_prime = victim_rsa.decrypt(tampered_ciphertext_c_prime)
    print(f"4. 受害者解密 C' 后，得到结果 M' = {decrypted_m_prime}")

    # --- 攻击者恢复原始明文 ---
    # 8. 攻击者计算 r 的模逆元
    r_inverse = mod_inverse(r, victim_rsa.n)
    
    # 9. 攻击者用 M' 和 r_inverse 恢复原始明文 M
    recovered_message_int = (decrypted_m_prime * r_inverse) % victim_rsa.n
    print(f"5. 攻击者计算 M = M' * r^-1 mod n 来恢复原始消息。")
    
    # --- 验证 ---
    recovered_message_bytes = recovered_message_int.to_bytes((recovered_message_int.bit_length() + 7) // 8, 'big')
    
    print("\n" + "="*20, "攻击结果", "="*20)
    print(f"攻击者恢复出的消息: '{recovered_message_bytes.decode()}'")
    
    if recovered_message_int == original_message_int:
        print("++ 攻击成功！攻击者在不知道私钥的情况下，完全恢复了原始的秘密消息。")
    else:
        print("-- 攻击失败，代码或逻辑有误。")



if __name__ == '__main__':
    # --- 运行CCA测试 ---
    print("\n\n" + "#"*10 + " Part 1: 选择密文攻击测试 " + "#"*10)
    test_rsa_cca_resistance()
    test_elgamal_malleability_cca()
    test_ecc_cca_resistance()
    demonstrate_textbook_rsa_cca()

    # --- 运行协议级和物理攻击模拟 ---
    print("\n\n" + "#"*10 + " Part 2: 其他攻击模拟 " + "#"*10)
    test_ecc_nonce_reuse_attack()
    test_rsa_common_modulus_attack()
    test_fault_injection_simulation()

    # --- 运行边信道攻击模拟 ---
    print("\n\n" + "#"*10 + " Part 3: 边信道攻击模拟 " + "#"*10)
    simulate_timing_attack_vulnerability()